import os
import argparse
import torch
import numpy as np
import cv2
from PIL import Image
import matplotlib.pyplot as plt

# ========== 导入 MedSAM2 模型构建 ==========
from medsam.build_medsam2 import build_medsam2  # 确保项目根目录有 medsam2/
from medsam.utils.inference import inference_one_image

# ========== 参数解析 ==========
def parse_args():
    parser = argparse.ArgumentParser(description="MedSAM2 Ultrasound Inference")
    parser.add_argument("--image", type=str, required=True, help="输入的超声图像路径 (.png/.jpg)")
    parser.add_argument("--model", type=str, required=True, help="MedSAM2 模型权重路径 (.pt)")
    parser.add_argument("--output", type=str, default="results", help="输出目录")
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
    return parser.parse_args()

# ========== 主流程 ==========
def main():
    args = parse_args()
    os.makedirs(args.output, exist_ok=True)

    # 1️⃣ 加载模型
    print(f"🔹 Loading MedSAM2 model from {args.model}")
    medsam2 = build_medsam2(model_path=args.model, device=args.device)
    medsam2.eval()

    # 2️⃣ 读取图像
    image_path = args.image
    image_name = os.path.basename(image_path)
    image = np.array(Image.open(image_path).convert("RGB"))

    # 3️⃣ 模型推理
    print("🔹 Running segmentation inference ...")
    mask = inference_one_image(medsam2, image, device=args.device)
    mask = (mask > 0.5).astype(np.uint8)  # 二值化

    # 4️⃣ 生成可视化叠加图
    overlay = image.copy()
    overlay[mask == 1] = [255, 0, 0]  # 红色区域
    blended = cv2.addWeighted(image, 0.7, overlay, 0.3, 0)

    # 5️⃣ 保存结果
    output_path = os.path.join(args.output, image_name.replace(".png", "_masked.png"))
    Image.fromarray(blended).save(output_path)
    print(f"✅ Saved overlay result to: {output_path}")

    # 可选显示结果
    plt.figure(figsize=(10,5))
    plt.subplot(1,2,1)
    plt.imshow(image)
    plt.title("Original")
    plt.axis("off")

    plt.subplot(1,2,2)
    plt.imshow(blended)
    plt.title("Overlay Mask")
    plt.axis("off")

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()
